<template>
  <div class="app-container">
    <!-- 画布组件 -->
    <div class="canvas-container">
      <SquareCanvas 
        ref="canvas"
        :width="canvasSize"
        :height="canvasSize"
        @draw="handleDrawEvent"
      />
    </div>
    
    <!-- 按钮组 -->
    <div class="button-group">
      <ActionButton 
        label="清除画布" 
        @click="clearCanvas" 
        class="button"
      />
      <ActionButton 
        label="识别数字" 
        @click="recognizeDigit" 
        class="button"
      />
    </div>
    
    <!-- ONNX模型识别结果 -->
    <div class="cnn-container">
      <h2>数字识别结果</h2>
      <div :class="['status', { disabled: !session || loading }]">
        {{ loading ? '模型加载中...' : '模型已就绪' }}
      </div>
      <div v-if="result" class="result-display">{{ result }}</div>
    </div>
  </div>
</template>

<script>
import SquareCanvas from './SquareCanvas.vue'
import ActionButton from './ActionButton.vue'
import * as ort from 'onnxruntime-web';

export default {
  components: {
    SquareCanvas,
    ActionButton
  },
  data() {
    return {
      canvasSize: 400,
      session: null,
      loading: false,
      result: null,
      drawingData: []
    }
  },
  async mounted() {
    await this.loadModel();
  },
  methods: {
    // 处理画布绘制事件
    handleDrawEvent(eventData) {
      this.drawingData.push(eventData);
    },
    
    // 清除画布
    clearCanvas() {
      this.$refs.canvas.clear();
      this.drawingData = [];
      this.result = null;
    },
    
    // 加载ONNX模型
    async loadModel() {
      try {
        this.loading = true;
        const modelUrl = new URL('../../assets/models/mnist_cnn.onnx', import.meta.url).href;
        const response = await fetch(modelUrl);
        const modelBuffer = await response.arrayBuffer();
        this.session = await ort.InferenceSession.create(modelBuffer);
        console.log('MNIST模型加载成功');
      } catch (error) {
        console.error('模型加载失败:', error);
      } finally {
        this.loading = false;
      }
    },
    
    // 识别数字
    async recognizeDigit() {
      if (!this.session) {
        console.error('模型会话未初始化');
        return;
      }
      
      try {
        this.loading = true;
        const canvas = this.$refs.canvas.$el;
        const imageData = this.prepareImageData(canvas);
        await this.runInference(imageData);
      } catch (error) {
        console.error('识别过程中出错:', error);
        this.result = `识别错误: ${error.message}`;
      } finally {
        this.loading = false;
      }
    },
    
    // 准备图像数据
    prepareImageData(canvas) {
      const ctx = canvas.getContext('2d');
      const tempCanvas = document.createElement('canvas');
      tempCanvas.width = 28;
      tempCanvas.height = 28;
      const tempCtx = tempCanvas.getContext('2d');
      
      // 将绘图缩放到28x28
      tempCtx.drawImage(canvas, 0, 0, 28, 28);
      const imageData = tempCtx.getImageData(0, 0, 28, 28);
      
      // 转换为灰度并归一化
      const data = new Float32Array(28 * 28);
      for (let i = 0; i < imageData.data.length; i += 4) {
        const r = imageData.data[i];
        const g = imageData.data[i + 1];
        const b = imageData.data[i + 2];
        const gray = 0.299 * r + 0.587 * g + 0.114 * b;
        data[i / 4] = (255 - gray) / 255.0; // 反转并归一化
      }
      
      return data;
    },
    
    // 运行推理
    async runInference(imageData) {
      const inputTensor = new ort.Tensor('float32', imageData, [1, 1, 28, 28]);
      const outputs = await this.session.run({ input: inputTensor });
      const outputArray = Array.from(outputs.output.data);
      const predictedClass = outputArray.indexOf(Math.max(...outputArray));
      const confidence = Math.max(...outputArray);
      
      this.result = `识别结果: 数字 ${predictedClass} (置信度: ${confidence.toFixed(4)})`;
    }
  }
}
</script>

<style scoped>
.app-container {
  display: flex;
  flex-direction: column;
  align-items: center;
  padding: 20px;
  font-family: Arial, sans-serif;
  max-width: 800px;
  margin: 0 auto;
}

.canvas-container {
  margin: 20px 0;
  border: 2px solid #ffd4ef;
  border-radius: 8px;
  box-shadow: 0 2px 10px rgba(0,0,0,0.1);
}

.button-group {
  display: flex;
  justify-content: center;
  gap: 20px;
  margin: 20px 0;
}

.cnn-container {
  width: 100%;
  padding: 20px;
  margin-top: 20px;
  border: 1px solid #ffd6d9;
  border-radius: 8px;
  background-color: #fbfbfb;
}

.cnn-container h2 {
  margin-top: 0;
  color: #66ccff;
  text-align: center;
}

.status {
  padding: 10px;
  background-color: #66ccff;
  color: #ffffff;
  border-radius: 4px;
  text-align: center;
  margin-bottom: 15px;
}

.status.disabled {
  background-color: #ffd6d9;
  color: #66ccff;
}

.result-display {
  padding: 15px;
  background-color: white;
  border-radius: 6px;
  font-size: 1.2rem;
  text-align: center;
  font-weight: bold;
  color: #66ccff;
  border: 1px solid #ffffff;
  background-color: #ffffff;
}
</style>